from typing import Tuple

import torch
import torch.nn as nn


class YoloLoss(nn.Module):
    def __init__(
        self,
        lambda_coord: float,
        lambda_noobj: float,
        bounding_box_num: int,
        class_num: int,
    ) -> None:
        super().__init__()
        self.lambda_coord = lambda_coord
        self.lambda_noobj = lambda_noobj
        self.bounding_box_num = bounding_box_num
        self.class_num = class_num

    def forward(self, y_hat: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
        # shape: [bs, gs, gs, 5 * bbn + class_num]
        y_true_area = y_true[:, :, :, : 5 * self.bounding_box_num].reshape(
            (
                y_true.shape[0],
                y_true.shape[1],
                y_true.shape[2],
                self.bounding_box_num,
                5,
            )
        )
        y_true_class = y_true[:, :, :, 5 * self.bounding_box_num :]
        y_hat_area = y_hat[:, :, :, : 5 * self.bounding_box_num].reshape(
            (y_hat.shape[0], y_hat.shape[1], y_hat.shape[2], self.bounding_box_num, 5)
        )
        y_hat_class = y_hat[:, :, :, 5 * self.bounding_box_num :]

        # shape: [bs, gs, gs, bbn]
        iou = self._cal_iou(y_hat_area, y_true_area)
        idx = iou.argmax(dim=3, keepdim=True)
        y_hat_area_res = torch.take_along_dim(y_hat_area, idx.unsqueeze(-1), 3).squeeze(
            3
        )
        y_true_area_res = y_true_area[..., 0, :]
        obj_flag = self._get_obj_flag(y_true_area_res)
        y_true_area_res[..., 4] *= iou.max(dim=3, keepdim=True)[0].squeeze(-1)
        # bounding box center loss
        center_loss = (
            (
                self.lambda_coord
                * obj_flag
                * (
                    (y_true_area_res[..., 0] - y_hat_area_res[..., 0]) ** 2
                    + (y_true_area_res[..., 1] - y_hat_area_res[..., 1]) ** 2
                )
            )
            .sum(dim=(1, 2))
            .mean()
        )
        bbx_loss = (
            (
                self.lambda_coord
                * obj_flag
                * (
                    (
                        torch.sqrt(y_true_area_res[..., 2])
                        - torch.sqrt(y_hat_area_res[..., 2])
                    )
                    ** 2
                    + (
                        torch.sqrt(y_true_area_res[..., 3])
                        - torch.sqrt(y_hat_area_res[..., 3])
                    )
                    ** 2
                )
            )
            .sum(dim=(1, 2))
            .mean()
        )
        confidence_loss = (
            (
                obj_flag * (y_true_area_res[..., 4] - y_hat_area_res[..., 4]) ** 2
                + self.lambda_noobj
                * (~obj_flag)
                * (y_true_area_res[..., 4] - y_hat_area_res[..., 4]) ** 2
            )
            .sum(dim=(1, 2))
            .mean()
        )

        classification_loss = (
            (obj_flag * ((y_true_class - y_hat_class) ** 2).sum(3))
            .sum(dim=(1, 2))
            .mean()
        )

        loss = center_loss + bbx_loss + confidence_loss + classification_loss

        return {
            "loss": loss,
            "center_loss": center_loss,
            "bbx_loss": bbx_loss,
            "confidence_loss": confidence_loss,
            "classification_loss": classification_loss,
        }

    def _cal_iou(
        self, y_hat_area: torch.Tensor, y_true_area: torch.Tensor
    ) -> torch.Tensor:
        # shape: [bs, gs, gs, bbn, 5]
        y_hat_x_min, y_hat_x_max, y_hat_y_min, y_hat_y_max = self._cal_coor(y_hat_area)

        y_true_x_min, y_true_x_max, y_true_y_min, y_true_y_max = self._cal_coor(
            y_true_area
        )

        # 计算交集的宽度和高度
        wi = torch.min(y_hat_x_max, y_true_x_max) - torch.max(y_hat_x_min, y_true_x_min)
        hi = torch.min(y_hat_y_max, y_true_y_max) - torch.max(y_hat_y_min, y_true_y_min)

        # 确保交集的宽度和高度非负
        wi = torch.clamp(wi, min=0)
        hi = torch.clamp(hi, min=0)

        # 计算交集和并集的面积
        intersection = wi * hi
        y_hat_area = (y_hat_x_max - y_hat_x_min) * (y_hat_y_max - y_hat_y_min)
        y_true_area = (y_true_x_max - y_true_x_min) * (y_true_y_max - y_true_y_min)
        union = y_hat_area + y_true_area - intersection

        # 计算IoU
        iou = intersection / (union + 1e-6)  # 添加一个小的epsilon防止除以0

        return iou

    @torch.no_grad
    def _cal_coor(
        self, area: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        # shape: [bs, gs, gs, bbn, 5]
        # NOTE: center label is scaled to [0, 1], from [0, 1/grid_size], when calculating IoU, it should be re-scaled.
        grid_size = area.shape[-3]
        center_x = area[..., 0] / grid_size
        center_y = area[..., 1] / grid_size
        width = area[..., 2]
        height = area[..., 3]
        x_min = center_x - width / 2
        x_max = center_x + width / 2
        y_min = center_y - height / 2
        y_max = center_y + height / 2
        return x_min, x_max, y_min, y_max

    @torch.no_grad
    def _get_obj_flag(self, y_true_area: torch.Tensor) -> torch.Tensor:
        return y_true_area[..., 4] > 0
